import sqlite3
from datetime import datetime
import pytz
from src.core.models import Trade, Order
from src.storage.sqlalchemy_storage import SQLAlchemyStorage
import os
from project_paths import PROJECT_ROOT

def migrate_trades(old_conn, storage):
    """迁移 trades 表数据"""
    print("开始迁移 trades 表...")
    old_cursor = old_conn.cursor()
    
    # 读取旧数据库中的 trades 表数据
    print("正在读取 trades 表数据...")
    old_cursor.execute('SELECT * FROM trades')
    old_trades = old_cursor.fetchall()
    print(f"从旧数据库读取到 {len(old_trades)} 条交易记录")

    # 获取列名
    print("正在获取 trades 表结构...")
    old_cursor.execute('PRAGMA table_info(trades)')
    columns = [column[1] for column in old_cursor.fetchall()]
    print(f"trades 表包含以下列: {columns}")

    # 将旧数据转换为 SQLAlchemy Trade 对象
    print("正在转换 trades 数据...")
    new_trades = []
    for i, old_trade in enumerate(old_trades, 1):
        try:
            trade_dict = dict(zip(columns, old_trade))
            
            # 转换时间字段
            for time_field in ['open_datetime', 'close_datetime']:
                if trade_dict[time_field]:
                    trade_dict[time_field] = datetime.fromisoformat(trade_dict[time_field])
                else:
                    trade_dict[time_field] = None

            # 转换数值字段
            for num_field in ['open_price', 'open_qty', 'close_price', 'close_qty', 
                            'point_profit_per_lot', 'total_points', 'gross_profit', 
                            'net_profit', 'commission', 'max_floating_profit_points', 
                            'max_floating_profit_dollars', 'max_floating_loss_points', 
                            'max_floating_loss_dollars', 'stop_loss_price', 
                            'take_profit_price']:
                if trade_dict[num_field] is not None:
                    trade_dict[num_field] = float(trade_dict[num_field])

            # 转换整数字段
            if trade_dict['nth_trade'] is not None:
                trade_dict['nth_trade'] = int(trade_dict['nth_trade'])

            # 创建新的 Trade 对象，不设置 index 字段
            new_trade = Trade(**trade_dict)
            new_trades.append(new_trade)
            if i % 100 == 0:  # 每处理100条记录打印一次进度
                print(f"已处理 {i} 条交易记录...")
        except Exception as e:
            print(f"处理第 {i} 条交易记录时出错: {str(e)}")
            print(f"问题数据: {old_trade}")
            continue

    # 保存到新的数据库
    print("正在保存 trades 数据到新数据库...")
    storage.save_trades(new_trades)
    print(f"Successfully migrated {len(new_trades)} trades to the new database.")

def migrate_orders(old_conn, storage):
    """迁移 orders 表数据"""
    print("开始迁移 orders 表...")
    old_cursor = old_conn.cursor()
    
    # 读取旧数据库中的 orders 表数据
    print("正在读取 orders 表数据...")
    old_cursor.execute('SELECT * FROM orders')
    old_orders = old_cursor.fetchall()
    print(f"从旧数据库读取到 {len(old_orders)} 条订单记录")

    # 获取列名
    print("正在获取 orders 表结构...")
    old_cursor.execute('PRAGMA table_info(orders)')
    columns = [column[1] for column in old_cursor.fetchall()]
    print(f"orders 表包含以下列: {columns}")

    # 将旧数据转换为 SQLAlchemy Order 对象
    print("正在转换 orders 数据...")
    new_orders = []
    for i, old_order in enumerate(old_orders, 1):
        try:
            order_dict = dict(zip(columns, old_order))
            
            # 转换时间字段
            for time_field in ['open_time', 'close_time', 'expiration']:
                if order_dict[time_field]:
                    order_dict[time_field] = datetime.fromisoformat(order_dict[time_field])
                else:
                    order_dict[time_field] = None

            # 转换数值字段
            for num_field in ['quantity', 'filled', 'quantity_left', 'limit_price', 
                            'stop_price', 'avg_fill_price', 'commission', 'multiplier']:
                if order_dict[num_field] is not None:
                    order_dict[num_field] = float(order_dict[num_field])

            # 创建新的 Order 对象，不设置 index 字段
            new_order = Order(**order_dict)
            new_orders.append(new_order)
            if i % 100 == 0:  # 每处理100条记录打印一次进度
                print(f"已处理 {i} 条订单记录...")
        except Exception as e:
            print(f"处理第 {i} 条订单记录时出错: {str(e)}")
            print(f"问题数据: {old_order}")
            continue

    # 保存到新的数据库
    print("正在保存 orders 数据到新数据库...")
    storage.save_orders(new_orders)
    print(f"Successfully migrated {len(new_orders)} orders to the new database.")

def main():
    # 创建新的存储实例
    storage = SQLAlchemyStorage(
        db_path=os.path.join(PROJECT_ROOT, 'db', 'trades.db')
    )

    # 连接旧数据库
    old_db_path = os.path.join(PROJECT_ROOT, 'db', 'trades.db.bak')
    if not os.path.exists(old_db_path):
        print(f"备份数据库文件不存在: {old_db_path}")
        return

    old_conn = sqlite3.connect(old_db_path)
    
    try:
        # 迁移数据
        migrate_trades(old_conn, storage)
        migrate_orders(old_conn, storage)
        print("数据迁移完成！")
    finally:
        old_conn.close()

if __name__ == '__main__':
    main() 
